import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np


class CustomizableBPNN(nn.Module):
    def __init__(self, input_size, output_size, hidden_layers, neurons_per_layer, activation_function):
        super(CustomizableBPNN, self).__init__()
        layer_sizes = [input_size] + [neurons_per_layer] * hidden_layers + [output_size]
        self.layers = nn.ModuleList(
            [nn.Linear(layer_sizes[i], layer_sizes[i + 1]) for i in range(len(layer_sizes) - 1)])
        self.activation_function = self.get_activation_function(activation_function)

    def forward(self, x):
        for layer in self.layers[:-1]:
            x = self.activation_function(layer(x))
        return self.layers[-1](x)

    def get_activation_function(self, activation_function):
        if activation_function == 'sigmoid':
            return nn.Sigmoid()
        elif activation_function == 'relu':
            return nn.ReLU()
        elif activation_function == 'tanh':
            return nn.Tanh()
        else:
            raise ValueError("Unsupported activation function")


if __name__ == '__main__':
    # 示例使用：
    # 自定义参数
    input_size = 2
    output_size = 1
    hidden_layers = 2
    neurons_per_layer = 3
    activation_function = 'sigmoid'  # 可选 'sigmoid', 'relu', 'tanh' 等

    # 创建BP神经网络
    bpnn = CustomizableBPNN(input_size, output_size, hidden_layers, neurons_per_layer, activation_function)

    # 构造训练数据
    X_train = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=np.float32)
    y_train = np.array([[0], [1], [1], [0]], dtype=np.float32)

    X_train = torch.tensor(X_train)
    y_train = torch.tensor(y_train)

    # 定义损失函数和优化器
    criterion = nn.MSELoss()
    optimizer = optim.SGD(bpnn.parameters(), lr=0.1)

    # 训练BP神经网络
    epochs = 10000
    for epoch in range(epochs):
        optimizer.zero_grad()
        outputs = bpnn(X_train)
        loss = criterion(outputs, y_train)
        loss.backward()
        optimizer.step()
        print(f"Epoch {epoch + 1}/{epochs}, Loss: {loss.item():.6f}")

    # 测试数据
    X_test = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=np.float32)
    X_test = torch.tensor(X_test)
    predictions = bpnn(X_test)
    print("Predictions:")
    print(predictions.detach().numpy())



